from scipy.stats import chi2_contingency, pearsonr
import pandas as pd
import numpy as np



def ana_data(path):
    # 1. 获取数据
    data = pd.read_csv(path)
    # 皮尔逊系数判断 连续值变量与标签之间的关系
    # 选择变量列 ：年龄，家与工作点的距离，月收入，曾工作的公司数量，薪资涨幅百分比，总工作年限，在公司工作的年限，自上次晋升以来的年限，与当前经理共事年限，在当前岗位工作年限，去年参加培训次数，（11列），
    feature_coulmns = ['Age', 'DistanceFromHome', 'MonthlyIncome', 'NumCompaniesWorked', 'PercentSalaryHike',
                       'TotalWorkingYears',
                       'YearsAtCompany', 'YearsSinceLastPromotion', 'YearsWithCurrManager', 'YearsInCurrentRole',
                       'TrainingTimesLastYear']
    # 选择标签列
    x_1 = data[feature_coulmns]
    y_1 = data['Attrition']
    # print(y_1.values)
    pearson_list = []

    for i in feature_coulmns:
        # print(x_1[i].values)
        # 皮尔逊系数
        corr_matrix = np.corrcoef(x_1[i].values, y_1.values)[0, 1]
        pearson_list.append(corr_matrix)
    # print(pearson_list)


    dict_data = dict(zip(feature_coulmns, pearson_list))
    x_1_list = []
    for i, j in dict_data.items():
        if abs(j) > 0.03:
            x_1_list.append(i)

    return x_1_list


def dm02(path):
    data = pd.read_csv(path)
    # 斯皮尔曼秩相关系数。 判断有序分类数据变量与标签之间的关系
    # 选择列：教育程度，工作环境满意度，工作投入度，工作级别，工作满意度，绩效等级，关系满意度，股票期权等级，生活工作平衡度（共10列），是否离职
    feature_coulmns = ['Education', 'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel', 'JobSatisfaction',
                       'PerformanceRating', 'RelationshipSatisfaction', 'StockOptionLevel', 'WorkLifeBalance']
    x_2 = data[feature_coulmns]
    y_2 = data['Attrition']

    df1 = pd.concat([x_2, y_2], axis=1)
    corr_matrix1 = df1.corr(method='spearman')

    # fig, axes = plt.subplots(figsize=(20, 10))
    # bar = axes.bar(corr_matrix1.columns, corr_matrix1['Attrition'].values.tolist(), width=0.5)
    # plt.tight_layout()
    # plt.ylabel('spearman')
    # plt.show()
    # plt.savefig('../data/fig/斯皮尔曼秩相关系数.png')

    dict_data = dict(zip(corr_matrix1.columns, corr_matrix1['Attrition']))
    x_2_list = []
    for i, j in dict_data.items():
        if abs(j) > 0.03:
            if i == 'Attrition':
                continue
            else:
                x_2_list.append(i)

    return x_2_list


def perform_chi_square_test(observed_data, significance_level=0.05):
    # 卡方检验
    chi2, p, dof, expected = chi2_contingency(observed_data)
    # 确定是否拒绝原假设
    reject_null = p < significance_level
    # 构建结果字典
    results = {
        "卡方统计量": chi2,
        "p值": p,
        "自由度": dof,
        "显著性水平": significance_level,
        "是否拒绝原假设": reject_null,
        "观测频数": observed_data,
        "期望频数": expected
    }
    return results


def observed_data():
    observed1 = np.array([[99, 9],
                          [156, 46],
                          [664, 123]])
    # print("=== 变量商务旅行情况与是否离职的关系 ===")
    results1 = perform_chi_square_test(observed1)
    # print(results1)

    observed2 = np.array([[13, 6],
                          [392, 70],
                          [100, 27],
                          [291, 49],
                          [56, 7]])
    # print("\n=== 变量教育领域与是否离职的关系 ===")
    results2 = perform_chi_square_test(observed2)
    # print(results2)

    observed3 = np.array([[378, 69],
                          [544, 109]
                          ])
    # print("\n=== 变量性别与是否离职的关系 ===")
    result3 = perform_chi_square_test(observed3)
    # print(result3)

    observed4 = np.array([[75, 5],
                          [93, 8],
                          [54, 2],
                          [180, 41],
                          [205, 42],
                          [34, 23]
                          ])
    # print("\n=== 变量工作角色与是否离职的关系 ===")
    result4 = perform_chi_square_test(observed4)
    # print(result4)

    observed5 = np.array([[216, 22],
                          [438, 62],
                          [268, 94]
                          ])
    # print("\n=== 变量婚姻状况与是否离职的关系 ===")
    result5 = perform_chi_square_test(observed5)
    # print(result5)

    observed6 = np.array([[714, 80],
                          [208, 98]
                          ])
    # print("\n=== 变量是否加班与是否离职的关系 ===")
    result6 = perform_chi_square_test(observed6)
    # print(result6)

    x_3_list = ['BusinessTravel', 'JobRole', 'MaritalStatus', 'OverTime']
    return x_3_list


def feature_engineering(path):
    data = pd.read_csv(path)
    # 获取特征和目标
    x1 = ana_data(path)
    x2 = dm02(path)
    x3 = observed_data()
    x_list = x1 + x2 + x3
    x = data[x_list]
    # print(train_x)
    y = data['Attrition']
    return x, y,x_list


if __name__ == '__main__':
    x, y ,x_list= feature_engineering('../new_project/data/train.csv')
    print(x_list)






